package com.zhang.graph.undirectedgraph;

import jdk.internal.org.objectweb.asm.tree.InnerClassNode;

import java.util.*;

/**
 * @author zhang
 * @time 2022/03/20 16:08:56
 */
public class WeightedGraph01 {
    private class Node {
        private String label;
        private List<Edge> edges = new ArrayList<>();

        public Node(String label) {
            this.label = label;
        }

        public void addEdge(Node to, int weight) {
            edges.add(new Edge(this, to, weight));
        }
        public List<Edge> getEdges() {
            return edges;
        }

        @Override
        public String toString() {
            return label;
        }
    }

    private class Edge {
        private Node from;
        private Node to;
        private int weight;

        public Edge(Node from, Node to, int weight) {
            this.from = from;
            this.to = to;
            this.weight = weight;
        }

        @Override
        public String toString() {
            return from + "->" + to;

        }
    }

    private Map<String, Node> nodes = new HashMap<>();

    public void addNode(String label) {
        Node node = new Node(label);
        nodes.putIfAbsent(label, node);
    }

    public void addEdge(String from, String to, int weight) {
        Node fromNode = nodes.get(from);
        if (fromNode == null) {
            return;
        }
        Node toNode = nodes.get(to);
        if (toNode == null) {
            return;
        }
        fromNode.addEdge(toNode,weight);
        toNode.addEdge(fromNode,weight);
    }

    public Path getShortestPath(String from, String to) {
        Node fromNode = nodes.get(from);
        Node toNode = nodes.get(to);

        HashMap<Node, Integer> distances = new HashMap<>();
        for (Node node : nodes.values()) {
            distances.put(node, Integer.MAX_VALUE);
        }
        distances.replace(fromNode, 0);

        HashMap<Node, Node> previous = new HashMap<>();
        HashSet<Node> visited = new HashSet<>();

        PriorityQueue<NodeEntry> queue = new PriorityQueue<>(Comparator.comparingInt(ne -> ne.priority));

        queue.add(new NodeEntry(fromNode, 0));
        while (!queue.isEmpty()) {
            Node current = queue.remove().node;
            visited.add(current);
            for (Edge edge : current.getEdges()) {
                if (visited.contains(edge.to)) {
                    continue;
                }
                int newDistance = distances.get(current) + edge.weight;
                if (newDistance < distances.get(edge.to)) {
                    distances.replace(edge.to, newDistance);
                    previous.put(edge.to, current);
                    queue.add(new NodeEntry(edge.to, newDistance));
                }

            }

        }
       return buildPath(previous, toNode);

    }

    private Path buildPath(HashMap<Node, Node> previousNode, Node toNode) {
        Stack<Node> stack = new Stack<>();
        stack.push(toNode);
        Node previous = previousNode.get(toNode);
        while (previous != null) {
            stack.push(previous);
            previous = previousNode.get(previous);
        }
        Path path = new Path();
        while (!stack.isEmpty()) {
            path.addNode(stack.pop().label);
        }
        return path;

    }

    private class NodeEntry {
        private Node node;
        private int priority;

        public NodeEntry(Node node, int priority) {
            this.node = node;
            this.priority = priority;
        }

        @Override
        public String toString() {
            return "NodeEntry{" +
                    "node=" + node +
                    ", priority=" + priority +
                    '}';
        }
    }
    public class Path {
        private List<String> nodes = new ArrayList<>();

        public void addNode(String node) {
            nodes.add(node);
        }
        @Override
        public String toString() {
            return nodes.toString();
        }
    }
    public WeightedGraph01 getMinimumSpanningTree() {
        WeightedGraph01 tree = new WeightedGraph01();

        if (nodes.isEmpty()) {
            return tree;
        }

        PriorityQueue<Edge> edges = new PriorityQueue<>(
                Comparator.comparingInt(e -> e.weight)
        );

        Node startNode = nodes.values().iterator().next();
        edges.addAll(startNode.getEdges());
        tree.addNode(startNode.label);

        while (tree.nodes.size() < nodes.size()) {
            Edge minEdge = edges.remove();
            Node nextNode = minEdge.to;

            if (tree.containsNode(nextNode.label)) {
                continue;
            }
            tree.addNode(nextNode.label);
            tree.addEdge(minEdge.from.label, nextNode.label, minEdge.weight);

            for (Edge edge : nextNode.getEdges()) {
                if (!tree.containsNode(edge.to.label)) {
                    edges.add(edge);
                }
            }

        }
        return tree;
    }

    public void print() {
        for (Node node : nodes.values()) {
            List<Edge> edges = node.getEdges();
            if (!edges.isEmpty()) {
                System.out.println(node + " is connected to " + edges);
            }

        }
    }

    public boolean containsNode(String label) {
        return nodes.containsKey(label);
    }

    public static void main(String[] args) {
        WeightedGraph01 wg = new WeightedGraph01();
        wg.addNode("A");
        wg.addNode("B");
        wg.addNode("C");
        wg.addNode("D");

        wg.addEdge("A","B",3);
        wg.addEdge("A","C",4);
        wg.addEdge("A","D",2);
        wg.addEdge("C","D",1);
        wg.addEdge("B","D",5);

        wg.getMinimumSpanningTree().print();

    }





}
